{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e217c4f2-4817-4ddf-a3c3-559ff22f632a",
   "metadata": {},
   "source": [
    "# Example 10 - Using the YOLO Annotation Tool\n",
    "This notebook showcases the YOLO annotation tool to draw boxes around Signals of Interest (SOIs) to create a image dataset compatible with YOLO.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53160dd6",
   "metadata": {},
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db8cb1a0-987c-448c-bbb3-2f7d3d908860",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchsig.image_datasets.annotation_tools.yolo_annotation_tool import yolo_annotator\n",
    "from torchsig.image_datasets.datasets.yolo_datasets import YOLOFileDataset\n",
    "from torchsig.image_datasets.plotting.plotting import plot_yolo_datum\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35b25503",
   "metadata": {},
   "source": [
    "## Set Directory Paths\n",
    "`unlabeled_image_dir/` is where the unlabelled spectrogram images are stored. After annotation, the annotated images with be written to `new_yolo_dataset_dir/`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a1d6ec5-31b0-4739-90ca-bcec3b5258c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "unlabeled_image_dir = \"./datasets/10_example_data/sample_images/\" # directory of images to be annotated\n",
    "new_yolo_dataset_dir = \"./datasets/10_example_data/annotated_yolo_sample_images/\" # directory to save annotated yolo data; be sure to end with a '/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "123b9d5b",
   "metadata": {},
   "source": [
    "## Create the YOLO Annotator and Label Signals\n",
    "Run the below command, and manually label the SOIs by drawing bounding boxes. Select which signal class you are annotating at the bottom, and hit `Submit` when done annotating the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f2ae761-b7f8-4f46-88c1-fb0c23c027cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolo_annotator(unlabeled_image_dir,new_yolo_dataset_dir, class_names=['signal_type_A', 'signal_type_B', 'signal_type_C'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cb224b3-45e3-4537-a248-43c5964441b6",
   "metadata": {},
   "source": [
    "## Plotting the Resulting Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3dcf767-0ffc-407a-bff2-a48e605f45c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "yds = YOLOFileDataset(new_yolo_dataset_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb54d59b-6358-4393-b15c-057810add6d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_yolo_datum(yds.next())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e353115c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
